"""
Contextual bandit algorithms for tool selection in Claude Code.

This module implements LinUCB and other contextual bandit algorithms as described
in Chapter 3 of the theoretical framework, providing regret bounds and optimal
tool selection strategies.
"""

using LinearAlgebra
using Random
using Distributions
using Statistics

# ============================================================================
# LinUCB Implementation
# ============================================================================

"""
    LinUCBAgent <: AbstractBandit

Linear contextual bandit using upper confidence bounds.
Implements Algorithm 3.1 from the theoretical framework with O(d√T log T) regret bounds.
"""
mutable struct LinUCBAgent <: AbstractBandit
    d::Int  # feature dimension
    k::Int  # number of arms (tools)
    α::Float64  # confidence parameter
    
    # Per-arm statistics
    A::Vector{Matrix{Float64}}  # covariance matrices A_i
    b::Vector{Vector{Float64}}  # reward vectors b_i
    θ::Vector{Vector{Float64}}  # parameter estimates θ̂_i
    
    # Exploration tracking
    total_rounds::Int
    arm_counts::Vector{Int}
    cumulative_reward::Float64
    regret_bound_cache::Vector{Float64}
    
    # Performance metrics
    confidence_widths::Vector{Vector{Float64}}
    prediction_errors::Vector{Float64}
    
    function LinUCBAgent(d::Int, k::Int, α::Float64 = 1.0, λ::Float64 = 1.0)
        # Initialize per-arm statistics
        A = [λ * I(d) for _ in 1:k]  # regularized identity matrices
        b = [zeros(d) for _ in 1:k]
        θ = [zeros(d) for _ in 1:k]
        
        new(d, k, α, A, b, θ, 0, zeros(Int, k), 0.0, Float64[], 
            Vector{Vector{Float64}}(), Float64[])
    end
end

"""
    select_arm(agent::LinUCBAgent, context::Vector{Float64}) -> Tuple{Int, Float64}

Select arm using LinUCB algorithm.
Returns (arm_index, confidence_score).
"""
function select_arm(agent::LinUCBAgent, context::Vector{Float64})
    @assert length(context) == agent.d "Context dimension mismatch"
    
    ucb_values = zeros(agent.k)
    confidence_widths = zeros(agent.k)
    
    for i in 1:agent.k
        # Compute UCB value: x^T θ̂_i + α√(x^T A_i^{-1} x)
        try
            A_inv = inv(agent.A[i])
            θ_est = A_inv * agent.b[i]
            confidence_width = agent.α * sqrt(dot(context, A_inv * context))
            
            ucb_values[i] = dot(context, θ_est) + confidence_width
            confidence_widths[i] = confidence_width
            
            # Update cached estimates
            agent.θ[i] = θ_est
            
        catch e
            @warn "Numerical instability in arm $i, using fallback" exception=e
            # Fallback to simple average with high uncertainty
            if agent.arm_counts[i] > 0
                avg_reward = agent.cumulative_reward / max(agent.total_rounds, 1)
                ucb_values[i] = avg_reward + agent.α * sqrt(2.0)
            else
                ucb_values[i] = agent.α * sqrt(2.0)  # High uncertainty for unplayed arms
            end
            confidence_widths[i] = agent.α * sqrt(2.0)
        end
    end
    
    # Select arm with highest UCB value
    selected_arm = argmax(ucb_values)
    confidence_score = confidence_widths[selected_arm]
    
    # Store confidence widths for analysis
    push!(agent.confidence_widths, copy(confidence_widths))
    
    return selected_arm, confidence_score
end

"""
    update_bandit!(agent::LinUCBAgent, arm::Int, reward::Float64, context::Vector{Float64})

Update LinUCB model with observed reward.
"""
function update_bandit!(agent::LinUCBAgent, arm::Int, reward::Float64, context::Vector{Float64})
    @assert 1 <= arm <= agent.k "Invalid arm index"
    @assert length(context) == agent.d "Context dimension mismatch"
    
    # Update statistics for selected arm
    agent.A[arm] += context * context'
    agent.b[arm] += reward * context
    
    # Update counters and tracking
    agent.total_rounds += 1
    agent.arm_counts[arm] += 1
    agent.cumulative_reward += reward
    
    # Compute prediction error for monitoring
    if agent.total_rounds > 1
        predicted_reward = dot(context, agent.θ[arm])
        prediction_error = abs(reward - predicted_reward)
        push!(agent.prediction_errors, prediction_error)
    end
end

"""
    get_regret_bound(agent::LinUCBAgent, horizon::Int, confidence::Float64 = 0.05) -> Float64

Compute theoretical regret bound for LinUCB.
Returns O(d√T log(kT/δ)) bound from Theorem 3.1.
"""
function get_regret_bound(agent::LinUCBAgent, horizon::Int, confidence::Float64 = 0.05)
    if horizon <= 0
        return 0.0
    end
    
    d, k = agent.d, agent.k
    δ = confidence
    
    # Theoretical bound: O(d√T log(kT/δ))
    log_term = log(k * horizon / δ)
    bound = agent.α * d * sqrt(horizon * log_term)
    
    # Cache for efficiency
    if length(agent.regret_bound_cache) < horizon
        resize!(agent.regret_bound_cache, horizon)
        for t in 1:horizon
            agent.regret_bound_cache[t] = agent.α * d * sqrt(t * log_term)
        end
    end
    
    return bound
end

"""
    compute_confidence_ellipsoid(agent::LinUCBAgent, arm::Int, confidence::Float64 = 0.95)

Compute confidence ellipsoid for arm parameters.
Returns (center, covariance_matrix).
"""
function compute_confidence_ellipsoid(agent::LinUCBAgent, arm::Int, confidence::Float64 = 0.95)
    @assert 1 <= arm <= agent.k "Invalid arm index"
    
    try
        A_inv = inv(agent.A[arm])
        center = A_inv * agent.b[arm]
        
        # Scale covariance by confidence level
        chi_squared_quantile = quantile(Chisq(agent.d), confidence)
        covariance = chi_squared_quantile * A_inv
        
        return center, covariance
        
    catch e
        @warn "Could not compute confidence ellipsoid for arm $arm" exception=e
        return zeros(agent.d), I(agent.d)
    end
end

"""
    get_arm_statistics(agent::LinUCBAgent, arm::Int) -> NamedTuple

Get detailed statistics for a specific arm.
"""
function get_arm_statistics(agent::LinUCBAgent, arm::Int)
    @assert 1 <= arm <= agent.k "Invalid arm index"
    
    try
        A_inv = inv(agent.A[arm])
        θ_est = A_inv * agent.b[arm]
        confidence_radius = agent.α * sqrt(tr(A_inv))
        
        return (
            parameter_estimate = θ_est,
            covariance_matrix = A_inv,
            confidence_radius = confidence_radius,
            num_selections = agent.arm_counts[arm],
            selection_frequency = agent.arm_counts[arm] / max(agent.total_rounds, 1)
        )
    catch e
        @warn "Could not compute statistics for arm $arm" exception=e
        return (
            parameter_estimate = zeros(agent.d),
            covariance_matrix = I(agent.d),
            confidence_radius = Inf,
            num_selections = agent.arm_counts[arm],
            selection_frequency = agent.arm_counts[arm] / max(agent.total_rounds, 1)
        )
    end
end

# ============================================================================
# Context Feature Engineering
# ============================================================================

"""
    ContextFeatureExtractor

Feature extraction for contextual bandit algorithms.
Converts high-level context into fixed-dimensional feature vectors.
"""
struct ContextFeatureExtractor
    feature_dim::Int
    feature_names::Vector{String}
    normalization_params::Dict{String, Tuple{Float64, Float64}}  # (mean, std)
    
    function ContextFeatureExtractor(feature_dim::Int)
        feature_names = [
            # Task characteristics
            "task_type_exploration", "task_type_implementation", "task_type_testing",
            "task_type_debugging", "task_type_refactoring", "task_type_documentation",
            
            # Codebase characteristics  
            "codebase_size_log", "num_files_log", "avg_file_size_log",
            "dependency_graph_density", "cyclomatic_complexity_avg",
            
            # Recent performance
            "recent_success_rate", "recent_avg_latency", "recent_error_rate",
            "tool_usage_frequency", "user_satisfaction_score",
            
            # Resource constraints
            "memory_pressure", "cpu_load", "disk_usage", "network_latency",
            "time_pressure", "budget_remaining",
            
            # Context history
            "conversation_length", "num_previous_tools", "similar_task_performance",
            "user_expertise_level", "code_quality_trend"
        ]
        
        # Pad or truncate to desired dimension
        if length(feature_names) < feature_dim
            append!(feature_names, ["feature_$(i)" for i in (length(feature_names)+1):feature_dim])
        elseif length(feature_names) > feature_dim
            feature_names = feature_names[1:feature_dim]
        end
        
        normalization_params = Dict{String, Tuple{Float64, Float64}}()
        
        new(feature_dim, feature_names, normalization_params)
    end
end

"""
    extract_features(extractor::ContextFeatureExtractor, context_state::ContextState,
                    task::CodeTask, environment::EnvironmentalState) -> Vector{Float64}

Extract feature vector from context.
"""
function extract_features(extractor::ContextFeatureExtractor, context_state::ContextState,
                         task::CodeTask, environment::EnvironmentalState)
    features = zeros(extractor.feature_dim)
    
    # Task type features (one-hot encoding)
    if extractor.feature_dim >= 6
        task_idx = if task.task_type == EXPLORATION
            1
        elseif task.task_type == IMPLEMENTATION
            2
        elseif task.task_type == TESTING
            3
        elseif task.task_type == DEBUGGING
            4
        elseif task.task_type == REFACTORING
            5
        else
            6  # DOCUMENTATION or other
        end
        features[task_idx] = 1.0
    end
    
    idx = 7  # Continue from task features
    
    # Codebase characteristics
    if idx <= extractor.feature_dim
        num_files = length(environment.filesystem.files)
        features[idx] = log1p(num_files)  # log(1 + num_files)
        idx += 1
    end
    
    if idx <= extractor.feature_dim
        total_size = sum(length(content) for content in values(environment.filesystem.files))
        avg_size = total_size / max(1, length(environment.filesystem.files))
        features[idx] = log1p(avg_size)
        idx += 1
    end
    
    if idx <= extractor.feature_dim
        num_nodes = length(environment.dependencies.nodes)
        num_edges = length(environment.dependencies.edges)
        density = num_edges / max(1, num_nodes * (num_nodes - 1))
        features[idx] = density
        idx += 1
    end
    
    # Recent performance features
    if idx <= extractor.feature_dim && !isempty(context_state.conversation_history.performance_history)
        recent_performance = context_state.conversation_history.performance_history[max(1, end-9):end]
        features[idx] = mean(recent_performance)
        idx += 1
    end
    
    # Resource features  
    if idx <= extractor.feature_dim
        features[idx] = environment.external.system_load / 10.0  # normalize to ~[0,1]
        idx += 1
    end
    
    if idx <= extractor.feature_dim
        memory_gb = environment.external.available_memory / 1e9
        features[idx] = min(1.0, memory_gb / 16.0)  # normalize to 16GB baseline
        idx += 1
    end
    
    # Context history features
    if idx <= extractor.feature_dim
        features[idx] = length(context_state.conversation_history.messages) / 100.0
        idx += 1
    end
    
    if idx <= extractor.feature_dim
        features[idx] = min(1.0, task.priority)
        idx += 1
    end
    
    # Fill remaining features with derived/interaction terms
    while idx <= extractor.feature_dim
        # Create interaction features or polynomial terms
        base_idx = ((idx - 1) % (min(idx-1, 5))) + 1
        if base_idx < idx
            features[idx] = features[base_idx]^2  # Polynomial features
        else
            features[idx] = 0.0  # Zero padding
        end
        idx += 1
    end
    
    # Normalize features
    return normalize_features(extractor, features)
end

"""
    normalize_features(extractor::ContextFeatureExtractor, features::Vector{Float64}) -> Vector{Float64}

Normalize feature vector using stored parameters.
"""
function normalize_features(extractor::ContextFeatureExtractor, features::Vector{Float64})
    normalized = copy(features)
    
    for (i, feature_name) in enumerate(extractor.feature_names)
        if haskey(extractor.normalization_params, feature_name)
            μ, σ = extractor.normalization_params[feature_name]
            if σ > 0
                normalized[i] = (features[i] - μ) / σ
            end
        end
    end
    
    # Clip to reasonable range to avoid numerical issues
    clamp!(normalized, -5.0, 5.0)
    
    return normalized
end

"""
    update_normalization!(extractor::ContextFeatureExtractor, features::Vector{Float64})

Update normalization parameters with new feature observation.
"""
function update_normalization!(extractor::ContextFeatureExtractor, features::Vector{Float64})
    α = 0.01  # Learning rate for exponential moving average
    
    for (i, feature_name) in enumerate(extractor.feature_names)
        feature_value = features[i]
        
        if haskey(extractor.normalization_params, feature_name)
            μ_old, σ_old = extractor.normalization_params[feature_name]
            
            # Update mean with exponential moving average
            μ_new = (1 - α) * μ_old + α * feature_value
            
            # Update standard deviation (simplified)
            σ_new = (1 - α) * σ_old + α * abs(feature_value - μ_new)
            σ_new = max(σ_new, 0.01)  # Prevent zero standard deviation
            
            extractor.normalization_params[feature_name] = (μ_new, σ_new)
        else
            # Initialize normalization parameters
            extractor.normalization_params[feature_name] = (feature_value, 1.0)
        end
    end
end

# ============================================================================
# Performance Analysis and Diagnostics
# ============================================================================

"""
    compute_empirical_regret(agent::LinUCBAgent, true_rewards::Vector{Vector{Float64}}, 
                           contexts::Vector{Vector{Float64}}, 
                           selected_arms::Vector{Int}) -> Vector{Float64}

Compute empirical regret given ground truth rewards.
"""
function compute_empirical_regret(agent::LinUCBAgent, true_rewards::Vector{Vector{Float64}},
                                contexts::Vector{Vector{Float64}}, 
                                selected_arms::Vector{Int})
    T = length(selected_arms)
    cumulative_regret = zeros(T)
    
    for t in 1:T
        context = contexts[t]
        selected_arm = selected_arms[t]
        
        # Compute optimal reward for this context
        optimal_reward = maximum(dot(context, true_rewards[i]) for i in 1:agent.k)
        
        # Compute achieved reward
        achieved_reward = dot(context, true_rewards[selected_arm])
        
        # Update cumulative regret
        instant_regret = optimal_reward - achieved_reward
        cumulative_regret[t] = (t > 1 ? cumulative_regret[t-1] : 0.0) + instant_regret
    end
    
    return cumulative_regret
end

"""
    analyze_arm_selection_pattern(agent::LinUCBAgent) -> NamedTuple

Analyze arm selection patterns and exploration behavior.
"""
function analyze_arm_selection_pattern(agent::LinUCBAgent)
    if agent.total_rounds == 0
        return (
            selection_entropy = 0.0,
            exploration_rate = 0.0,
            arm_diversity = 0.0,
            confidence_trend = Float64[]
        )
    end
    
    # Selection entropy
    selection_probs = agent.arm_counts / agent.total_rounds
    selection_entropy = -sum(p * log(p + 1e-10) for p in selection_probs if p > 0)
    
    # Exploration rate (fraction of arms tried)
    arms_tried = sum(agent.arm_counts .> 0)
    exploration_rate = arms_tried / agent.k
    
    # Arm diversity (inverse of concentration)
    arm_diversity = 1.0 - maximum(selection_probs)
    
    # Confidence trend (average confidence width over time)
    confidence_trend = if !isempty(agent.confidence_widths)
        [mean(widths) for widths in agent.confidence_widths]
    else
        Float64[]
    end
    
    return (
        selection_entropy = selection_entropy,
        exploration_rate = exploration_rate, 
        arm_diversity = arm_diversity,
        confidence_trend = confidence_trend,
        prediction_error_trend = copy(agent.prediction_errors)
    )
end

"""
    export_agent_state(agent::LinUCBAgent) -> Dict{String, Any}

Export agent state for serialization/checkpointing.
"""
function export_agent_state(agent::LinUCBAgent)
    return Dict{String, Any}(
        "d" => agent.d,
        "k" => agent.k,
        "α" => agent.α,
        "A" => [Matrix(A) for A in agent.A],  # Convert to regular matrices
        "b" => [Vector(b) for b in agent.b],
        "θ" => [Vector(θ) for θ in agent.θ],
        "total_rounds" => agent.total_rounds,
        "arm_counts" => copy(agent.arm_counts),
        "cumulative_reward" => agent.cumulative_reward,
        "confidence_widths" => deepcopy(agent.confidence_widths),
        "prediction_errors" => copy(agent.prediction_errors)
    )
end

"""
    import_agent_state!(agent::LinUCBAgent, state::Dict{String, Any})

Import agent state from serialized data.
"""
function import_agent_state!(agent::LinUCBAgent, state::Dict{String, Any})
    # Verify compatibility
    if state["d"] != agent.d || state["k"] != agent.k
        error("Agent dimensions don't match imported state")
    end
    
    agent.α = state["α"]
    agent.A = [Matrix{Float64}(A) for A in state["A"]]
    agent.b = [Vector{Float64}(b) for b in state["b"]]
    agent.θ = [Vector{Float64}(θ) for θ in state["θ"]]
    agent.total_rounds = state["total_rounds"]
    agent.arm_counts = Vector{Int}(state["arm_counts"])
    agent.cumulative_reward = state["cumulative_reward"]
    agent.confidence_widths = deepcopy(state["confidence_widths"])
    agent.prediction_errors = Vector{Float64}(state["prediction_errors"])
    
    return agent
end